#!/usr/bin/env python3
"""
服务器初始化模块
负责命令行参数解析、模型加载和管理器实例创建
"""

import argparse
from typing import Any

# 导入配置
from speech_convert.config.funasr_config import CONFIG
from speech_convert.server.response_manager import ResponseManager
from speech_convert.utils.console_style import (print_loading, print_stage,
                                                print_success)
from speech_convert.utils.model_manager import ModelManager
# 导入管理器模块
from speech_convert.utils.speaker_manager import SpeakerManager


class ServerInitializer:
    """服务器初始化器 - 负责参数解析和模型加载"""

    def __init__(self):
        self.args: argparse.Namespace | None = None
        self.websocket_users: set[Any] = set()
        self.speaker_manager: SpeakerManager | None = None
        self.response_manager: ResponseManager | None = None
        self.models: dict[str, Any] = {}

    def setup_argument_parser(self) -> argparse.ArgumentParser:
        """设置命令行参数解析器"""
        parser = argparse.ArgumentParser()

        # 服务器地址
        parser.add_argument(
            "--host",
            type=str,
            default=CONFIG.server.DEFAULT_HOST,
            required=False,
            help="host ip, localhost, 0.0.0.0",
        )
        # 服务器端口
        parser.add_argument(
            "--port",
            type=int,
            default=CONFIG.server.DEFAULT_PORT,
            required=False,
            help="grpc server port",
        )
        # vad 语音活动检测模型
        parser.add_argument(
            "--vad_model",
            type=str,
            default=CONFIG.model.VAD_MODEL,
            help="model from modelscope",
        )
        # vad 语音活动检测模型版本
        parser.add_argument(
            "--vad_model_revision",
            type=str,
            default=CONFIG.model.VAD_MODEL_REVISION,
            help="",
        )

        # asr 语音识别在线模型
        parser.add_argument(
            "--asr_model_online",
            type=str,
            default=CONFIG.model.ASR_MODEL_ONLINE,
            help="model from modelscope",
        )
        # asr 语音识别在线模型版本
        parser.add_argument(
            "--asr_model_online_revision",
            type=str,
            default=CONFIG.model.ASR_MODEL_ONLINE_REVISION,
            help="",
        )

        # asr 语音识别离线模型
        parser.add_argument(
            "--asr_model",
            type=str,
            default=CONFIG.model.ASR_MODEL,
            help="model from modelscope",
        )
        # asr 语音识别离线模型版本
        parser.add_argument(
            "--asr_model_revision",
            type=str,
            default=CONFIG.model.ASR_MODEL_REVISION,
            help="",
        )

        # 标点符号模型
        parser.add_argument(
            "--punc_model",
            type=str,
            default=CONFIG.model.PUNC_MODEL,
            help="model from modelscope",
        )
        # 标点符号模型版本
        parser.add_argument(
            "--punc_model_revision",
            type=str,
            default=CONFIG.model.PUNC_MODEL_REVISION,
            help="",
        )
        # GPU设置
        parser.add_argument(
            "--ngpu",
            type=int,
            default=CONFIG.server.DEFAULT_NGPU,
            help="0 for cpu, 1 for gpu",
        )
        parser.add_argument(
            "--device", type=str, default=CONFIG.server.DEFAULT_DEVICE, help="cuda, cpu"
        )
        parser.add_argument(
            "--ncpu", type=int, default=CONFIG.server.DEFAULT_NCPU, help="cpu cores"
        )

        # SSL证书设置
        parser.add_argument(
            "--certfile",
            type=str,
            default=CONFIG.server.DEFAULT_CERTFILE,
            required=False,
            help="certfile for ssl",
        )
        parser.add_argument(
            "--keyfile",
            type=str,
            default=CONFIG.server.DEFAULT_KEYFILE,
            required=False,
            help="keyfile for ssl",
        )

        # 日志级别
        parser.add_argument(
            "--log-level",
            type=str,
            default="info",
            required=False,
            help="日志级别 (info, debug, warning, error)"
        )

        return parser

    def parse_arguments(self) -> argparse.Namespace:
        """解析命令行参数"""
        parser = self.setup_argument_parser()
        self.args = parser.parse_args()
        return self.args

    def initialize_managers(self) -> tuple[SpeakerManager, ResponseManager]:
        """初始化管理器实例"""
        self.speaker_manager = SpeakerManager()
        self.response_manager = ResponseManager(self.speaker_manager)
        return self.speaker_manager, self.response_manager

    def load_models(self) -> ModelManager:
        """加载所有AI模型并返回ModelManager实例"""
        if not self.args:
            raise ValueError("请先调用 parse_arguments() 解析参数")

        print_stage("模型加载与验证")

        from funasr import AutoModel  # type: ignore

        # vad 语音活动检测
        print_loading("加载 VAD模型 ...")
        self.models["vad"] = AutoModel(
            model=self.args.vad_model,
            model_revision=self.args.vad_model_revision,
            ngpu=self.args.ngpu,
            ncpu=self.args.ncpu,
            device=self.args.device,
            disable_pbar=CONFIG.log.DISABLE_MODEL_PBAR,
            disable_log=CONFIG.log.DISABLE_MODEL_LOG,
            disable_update=CONFIG.log.DISABLE_UPDATE_CHECK,
        )
        print_success("VAD模型 加载完成")
        # asr 语音识别在线
        print_loading("加载 ASR在线模型 ...")
        self.models["asr_streaming"] = AutoModel(
            model=self.args.asr_model_online,
            model_revision=self.args.asr_model_online_revision,
            ngpu=self.args.ngpu,
            ncpu=self.args.ncpu,
            device=self.args.device,
            disable_pbar=CONFIG.log.DISABLE_MODEL_PBAR,
            disable_log=CONFIG.log.DISABLE_MODEL_LOG,
            disable_update=CONFIG.log.DISABLE_UPDATE_CHECK,
        )
        print_success("ASR在线模型 加载完成")
        # asr 语音识别离线
        print_loading("加载 ASR离线模型 ...")
        self.models["asr"] = AutoModel(
            model=self.args.asr_model,
            model_revision=self.args.asr_model_revision,
            ngpu=self.args.ngpu,
            ncpu=self.args.ncpu,
            device=self.args.device,
            disable_pbar=CONFIG.log.DISABLE_MODEL_PBAR,
            disable_log=CONFIG.log.DISABLE_MODEL_LOG,
            disable_update=CONFIG.log.DISABLE_UPDATE_CHECK,
        )
        print_success("ASR离线模型 加载完成")
        # 标点符号模型
        if self.args.punc_model != "":
            print_loading("加载 标点模型 ...")
            self.models["punc"] = AutoModel(
                model=self.args.punc_model,
                model_revision=self.args.punc_model_revision,
                ngpu=self.args.ngpu,
                ncpu=self.args.ncpu,
                device=self.args.device,
                disable_pbar=CONFIG.log.DISABLE_MODEL_PBAR,
                disable_log=CONFIG.log.DISABLE_MODEL_LOG,
                disable_update=CONFIG.log.DISABLE_UPDATE_CHECK,
            )
            print_success("标点模型 加载完成")
        else:
            self.models["punc"] = None

        # 创建ModelManager实例
        model_manager = ModelManager(self.models)

        return model_manager

    def initialize_all(
        self,
    ) -> tuple[
        argparse.Namespace, SpeakerManager, ResponseManager, ModelManager, set[Any]
    ]:
        """完整初始化 - 一次性完成所有初始化工作"""
        args = self.parse_arguments()
        speaker_manager, response_manager = self.initialize_managers()
        model_manager = self.load_models()

        return (
            args,
            speaker_manager,
            response_manager,
            model_manager,
            self.websocket_users,
        )


# 便捷函数
def initialize_server() -> tuple[
    argparse.Namespace, SpeakerManager, ResponseManager, ModelManager, set[Any]
]:
    """便捷函数 - 快速初始化整个服务器"""
    initializer = ServerInitializer()
    return initializer.initialize_all()
